#!/usr/bin/env octave

% Copyright (C) 2019 Kei Kebreau
%
% This program is free software: you can redistribute it and/or modify
% it under the terms of the GNU General Public License as published by
% the Free Software Foundation, either version 3 of the License, or
% (at your option) any later version.
%
% This program is distributed in the hope that it will be useful,
% but WITHOUT ANY WARRANTY; without even the implied warranty of
% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
% GNU General Public License for more details.
%
% You should have received a copy of the GNU General Public License
% along with this program.  If not, see <http://www.gnu.org/licenses/>.

clear -global;
% Let user select input file.
[input_file, file_path] = uigetfile({'*.xyz*', 'Supported file formats'});
% Load the input file.
raw_data = fileread([file_path input_file]);
[~, file_name, ~] = fileparts(input_file);
% Let the user input the desired resolution and distance from structure.
resolution = inputdlg('# of faces surrounding the sphere', 'Resolution');
resolution = str2num(resolution{1});
radius = inputdlg('Distance from structure', 'Radius');
radius = str2num(radius{1});

cd_regexp = ['Cd' repmat('\s+([+-]?[0-9]*\.?[0-9]+)', [1,3])];
cd_locations = [regexp(raw_data, cd_regexp, 'tokens'){:}];
cd_locations = reshape(sscanf(sprintf('%s ', cd_locations{:}), '%f'),
                       3, numel(cd_locations) / 3);
se_regexp = ['Se' repmat('\s+([+-]?[0-9]*\.?[0-9]+)', [1,3])];
se_locations = [regexp(raw_data, se_regexp, 'tokens'){:}];
se_locations = reshape(sscanf(sprintf('%s ', se_locations{:}), '%f'),
                       3, numel(se_locations) / 3);

% Get bounds for the grid.
[x_min, x_max] = bounds([cd_locations(1,:), se_locations(1,:)]);
x_mid = mean([x_min, x_max]);
[y_min, y_max] = bounds([cd_locations(2,:), se_locations(2,:)]);
y_mid = mean([y_min, y_max]);
[z_min, z_max] = bounds([cd_locations(3,:), se_locations(3,:)]);
z_mid = mean([z_min, z_max]);
base_radius = max([x_max - x_mid, y_max - y_mid, z_max - z_mid]) * 2;
radius += base_radius;
radii = 0.1:0.1:radius;
spheres = cell(1, length(radii));
loading_bar = waitbar(0, 'Generating spheres...');
for s = 1:length(spheres)
  tic;
  [X, Y, Z] = sphere(resolution);
  bench(s) = toc;
  X = radii(s) * X + x_mid;
  Y = radii(s) * Y + y_mid;
  Z = radii(s) * Z + z_mid;
  spheres{s} = {X, Y, Z};
  waitbar(s / length(spheres), loading_bar);
endfor

function [cd_distances, se_distances, potentials] = calc_potential(cd_locations,
                                                                   se_locations,
                                                                   X, Y, Z)
  potentials = zeros(size(X));
  q_e = 1.6021766208e-19;
  eps0 = 8.854187817e-12;
  kC = 1/(4*pi*eps0);
  cd_charge = 2 * q_e;
  se_charge = -2 * q_e;
  cd_distances = zeros(numel(X), length(cd_locations));
  se_distances = zeros(numel(X), length(se_locations));
  for ii = 1:length(cd_locations)
    potentials += cd_charge ./ hypot(cd_locations(1,ii) - X,
                                     cd_locations(2,ii) - Y,
                                     cd_locations(3,ii) - Z);
  endfor
  for ii = 1:length(se_locations)
    potentials += se_charge ./ hypot(se_locations(1,ii) - X,
                                     se_locations(2,ii) - Y,
                                     se_locations(3,ii) - Z);
  endfor
  potentials *= kC;
  potentials(potentials > 1) = NaN;
endfunction

cd_distances = se_distances = potentials = cell(1, length(spheres));
waitbar(0, loading_bar, 'Calculating electrostatic potential on each sphere...');
for s = 1:length(spheres)
  [cd_distances{s}, se_distances{s}, potentials{s}] = ...
    calc_potential(cd_locations, se_locations,
                   spheres{s}{1}, spheres{s}{2}, spheres{s}{3});
  waitbar(s / length(spheres), loading_bar);
endfor
close(loading_bar);

global g_potentials = potentials;

% Make a new layer of potential calculations visible.
function redraw_potential (handle, event)
  val = round(get(handle, 'value'));
  set(handle, 'value', val);
  global surfaces;
  global g_potentials;
  p = [g_potentials{val}](:);
  p = p(~isnan(p));
  caxis([min(p), max(p)]);
  cellfun(@(s) set(s, 'visible', 'off'), surfaces);
  set(surfaces{val}, 'visible', 'on');
  drawnow();
endfunction

% Takes "cd_locations" and "se_locations" matrices, and "spheres" and
% "potentials" cell arrays.
function f = plot_potential (cd_locations, se_locations, spheres, potentials)
  f = figure('Units', 'normalized',
             'Position', [0.1, 0.1, 0.7, 0.7],
             'visible', 'off');
  hold on;
  xlabel('X');
  ylabel('Y');
  zlabel('Z');
  title('Electrostatic potential (vastly oversimplified calculation)');
  scatter3(cd_locations(1,:), cd_locations(2,:), cd_locations(3,:),
           100, [0, 1, 0], '^', 'filled');
  scatter3(se_locations(1,:), se_locations(2,:), se_locations(3,:),
           100, [1, 0, 0], 'filled');
  global surfaces = cell(1, length(spheres));
  loading_bar = waitbar(0, 'Plotting spheres...');
  for ii = 1:length(surfaces)
    surfaces{ii} = surf(spheres{ii}{1},
                        spheres{ii}{2},
                        spheres{ii}{3},
                        potentials{ii});
    set(surfaces{ii}, 'visible', 'off', 'facealpha', 0.5);
    waitbar(ii / length(surfaces), loading_bar);
  endfor
  set(surfaces{end}, 'visible', 'on');
  colormap(flipud(jet));
  colorbar;
  p = [potentials{end}](:);
  p = p(~isnan(p));
  caxis([min(p), max(p)]);

  axis('auto', 'equal');
  rotate3d on;
  grid;
  view(3);
  hold off;
  hslider = uicontrol('parent', f,
                      'style', 'slider',
                      'Units', 'normalized',
                      'position', [0.1, 0.025, 0.8, 0.025],
                      'min', 1,
                      'max', length(spheres),
                      'value', length(spheres),
                      'sliderstep', [1/length(spheres), 10/length(spheres)],
                      'callback', @redraw_potential);
  close(loading_bar);
  set(f, 'visible', 'on');
endfunction

fig = plot_potential(cd_locations, se_locations, spheres, potentials);
waitfor(fig);
